/*
 * Copyright The OpenTelemetry Authors
 * SPDX-License-Identifier: Apache-2.0
 */

package com.xiashitech.agent.instrumentation.dubbo.common.inst;


import com.alibaba.dubbo.remoting.exchange.ResponseCallback;
import com.alibaba.dubbo.rpc.*;
import com.alibaba.dubbo.rpc.protocol.dubbo.FutureAdapter;
import com.xiashitech.agent.cache.status.CacheStatus;
import com.xiashitech.agent.instrumentation.utils.BaggageUtil;
import io.opentelemetry.javaagent.shaded.instrumentation.api.instrumenter.Instrumenter;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Context;
import io.opentelemetry.javaagent.shaded.io.opentelemetry.context.Scope;

import java.util.concurrent.Future;

final class TracingFilter implements Filter {

  private final Instrumenter<DubboRequest, Result> serverInstrumenter;
  private final Instrumenter<DubboRequest, Result> clientInstrumenter;

  TracingFilter(
      Instrumenter<DubboRequest, Result> serverInstrumenter,
      Instrumenter<DubboRequest, Result> clientInstrumenter) {
    this.serverInstrumenter = serverInstrumenter;
    this.clientInstrumenter = clientInstrumenter;
  }

  @Override
  public Result invoke(Invoker<?> invoker, Invocation invocation) {
    if(CacheStatus.cannotRun()) {
      return invoker.invoke(invocation);
    }
    if (!(invocation instanceof RpcInvocation)) {
      return invoker.invoke(invocation);
    }

    RpcContext rpcContext = RpcContext.getContext();
    boolean isServer = rpcContext.isProviderSide();
    Instrumenter<DubboRequest, Result> instrumenter =
            isServer ? serverInstrumenter : clientInstrumenter;
    Context parentContext = Context.current();
    DubboRequest request = DubboRequest.create((RpcInvocation) invocation, rpcContext);

    if (!instrumenter.shouldStart(parentContext, request)) {
      return invoker.invoke(invocation);
    }
    parentContext = BaggageUtil.syncBaggage().build().storeInContext(parentContext);
    Context context = instrumenter.start(parentContext, request);
    Context invokeContext = BaggageUtil.syncBaggage().build().storeInContext(context);
    Result result;
    boolean isSynchronous = true;
    try (Scope ignored = invokeContext.makeCurrent()) {
      result = invoker.invoke(invocation);
      if (!isServer) {
        Future<Object> future = rpcContext.getFuture();
        if (future != null) {
          isSynchronous = false;
          ((FutureAdapter) future).getFuture().setCallback(new ResponseCallback() {

            @Override
            public void done(Object o) {
              instrumenter.end(invokeContext, request, result, null);
            }

            @Override
            public void caught(Throwable throwable) {
              instrumenter.end(invokeContext, request, result, throwable);
            }
          });
        }
      }
    } catch (Throwable e) {
      instrumenter.end(invokeContext, request, null, e);
      throw e;
    }
    if (isSynchronous) {
      instrumenter.end(invokeContext, request, result, result.getException());
    }
    return result;
  }
}
